package temp;


import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.SecureRandom;
import java.util.Arrays;
import java.util.Base64;


/**
 * This is an implementation of the AES algorithm, specifically CBC mode,
 * with 256 bits key length and PKCS7 padding.
 * <p>
 * https://github.com/mervick/aes-everywhere/blob/master/java/src/com/github/mervick/aes_everywhere/Aes256.java
 */
public class AesCbcEverywhereManager {

    protected static final byte[] SALTED = "Salted__".getBytes(StandardCharsets.US_ASCII);

    /**
     * 加密方法
     *
     * @param input      待加密数据
     * @param passphrase 密码短语
     * @return Encrypted 加密结果
     * @throws Exception Throws exceptions
     */
    public byte[] encrypt(byte[] input, byte[] passphrase) throws Exception {
        byte[] salt = (new SecureRandom()).generateSeed(8);
        Object[] keyIv = deriveKeyAndIv(passphrase, salt);

        Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
        cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec((byte[]) keyIv[0], "AES"), new IvParameterSpec((byte[]) keyIv[1]));

        byte[] enc = cipher.doFinal(input);
        return concat(concat(SALTED, salt), enc);
    }


    /**
     * 解密方法
     *
     * @param data       待解密数据
     * @param passphrase 密码短语
     * @return Decrypted 解密结果
     * @throws Exception Throws exceptions
     */
    public byte[] decrypt(byte[] data, byte[] passphrase) throws Exception {
        byte[] salt = Arrays.copyOfRange(data, 8, 16);

        if (!Arrays.equals(Arrays.copyOfRange(data, 0, 8), SALTED)) {
            throw new IllegalArgumentException("Invalid crypted data");
        }

        Object[] keyIv = deriveKeyAndIv(passphrase, salt);

        Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
        cipher.init(Cipher.DECRYPT_MODE,
                new SecretKeySpec((byte[]) keyIv[0], "AES"),
                new IvParameterSpec((byte[]) keyIv[1]));
        return cipher.doFinal(data, 16, data.length - 16);
    }

    /**
     * 生成 key and iv
     *
     * @param passphrase 密码短语，即每个app生成的key
     * @param salt       盐
     * @return Array of key and iv
     * @throws Exception Throws exceptions
     */
    protected Object[] deriveKeyAndIv(byte[] passphrase, byte[] salt) throws Exception {
        final MessageDigest md5 = MessageDigest.getInstance("MD5");
        final byte[] passSalt = concat(passphrase, salt);
        byte[] dx = new byte[0];
        byte[] di = new byte[0];

        for (int i = 0; i < 3; i++) {
            di = md5.digest(concat(di, passSalt));
            dx = concat(dx, di);
        }

        return new Object[]{Arrays.copyOfRange(dx, 0, 32), Arrays.copyOfRange(dx, 32, 48)};
    }

    /**
     * 连接 bytes
     *
     * @param a First array
     * @param b Second array
     * @return Concatenated bytes
     */
    protected byte[] concat(byte[] a, byte[] b) {
        byte[] c = new byte[a.length + b.length];
        System.arraycopy(a, 0, c, 0, a.length);
        System.arraycopy(b, 0, c, a.length, b.length);
        return c;
    }

    /**
     * 测试
     */
    public static void main(String[] args) throws Exception {
        String plainText = "This is a plain text which need to be encrypted by Java AES 256 CBC Encryption Algorithm";
        String key = "itFJGm9ybahGSDdnn8lbi3Vr/jB3lhsMNCE3MvA9QjA=";

        AesCbcEverywhereManager aesCbcEverywhereManager = new AesCbcEverywhereManager();
        //byte[] encryptData = aesCbcEverywhereManager.encrypt(plainText.getBytes(StandardCharsets.UTF_8), Base64.getDecoder().decode(key));
        //System.out.println("Encrypted Text: " + Base64.getEncoder().encodeToString(encryptData));
        String data = "U2FsdGVkX1+PGrGCLvsTJBcqr5JLu/mLOsRdUWOC+OfT12dgGojncyZR/5H5GJIH1qlY4kJ3FK32q9GbkTAt4bD6BmlqWZ/DjXVLV9/MGCz8KsD+WYR5+xAJ8g0uf/Gh8okwrTe+FBs+ONymLlq4VyMYrDYBCM7uz+4ukx58V76cWhVMj6x9qrH1oNMtO5cVMwPXn/M6chGkaknPxyknvOpwaUmAwjV+rUy5mTDQUKESqvv9mcro+fQDkWoq56xgjrsirSgQyFP6zkMNRm+J3Dy6MBPRmtnEeYuzF89kbVAWI6P2bNE07JqkSXb/CFdPa/9af7RrLwnNwg31W/ke/3GwNUUaB5tmpslCwN0AE8BwrMX9uyXW11MaZcTffiQWl50Z8CVtiZ5hrJlouMKZumXlgGiWzaIsUGmMke9z/wsCcgvSkC2wM5M7nF+LQXRkczFgWQIKTPTSuxNsp3mzOA==";
        byte[] encryptData = Base64.getDecoder().decode(data);
        byte[] decryptData = aesCbcEverywhereManager.decrypt(encryptData, Base64.getDecoder().decode(key));
        System.out.println("Decrypted Text: " + new String(decryptData, StandardCharsets.UTF_8));
    }
}
